(* Mathematica Package *)

BeginPackage["Hamiltonian`",{"MathMR`", "MRData`","NVInteraction`"}]
(* Exported symbols added here with SymbolName::usage *)  
SPIN1x=SpinOperatorX[1];
SPIN1y=SpinOperatorY[1];
SPIN1z=SpinOperatorZ[1];
SPINhalfx=X/2;
SPINhalfy=Y/2;
SPINhalfz=Z/2;

SxN14 = KroneckerProduct[SPIN1x, IdentityMatrix[3]];
SyN14 = KroneckerProduct[SPIN1y, IdentityMatrix[3]];
SzN14 = KroneckerProduct[SPIN1z, IdentityMatrix[3]];
SvecN14 = {SxN14, SyN14, SzN14};
IxN14 = KroneckerProduct[IdentityMatrix[3], SPIN1x];
IyN14 = KroneckerProduct[IdentityMatrix[3], SPIN1y];
IzN14 = KroneckerProduct[IdentityMatrix[3], SPIN1z];
IvecN14 = {IxN14, IyN14, IzN14};

SxN15 = KroneckerProduct[SPIN1x, IdentityMatrix[2]];
SyN15 = KroneckerProduct[SPIN1y, IdentityMatrix[2]];
SzN15 = KroneckerProduct[SPIN1z, IdentityMatrix[2]];
SvecN15 = {SxN15, SyN15, SzN15};
IxN15 = KroneckerProduct[IdentityMatrix[3], SPINhalfx];
IyN15 = KroneckerProduct[IdentityMatrix[3], SPINhalfy];
IzN15 = KroneckerProduct[IdentityMatrix[3], SPINhalfz];
IvecN15 = {IxN15, IyN15, IzN15};

NVelectronSpin::usage="Hamiltonian of NV electron spin";
NVn14Spin::usage="Hamiltonian of NV N-14 nuclear spin";
NVn15Spin::usage="Hamiltonian of NV N-15 nuclear spin";
NV14hf::usage="hf interaction of NV electron spin and N-14 nuclear spin";
NV15hf::usage="hf interaction of NV electron spin and N-15 nuclear spin";

PrimeBasis::usage="Eigenstates of spin-1 operators along the given direction.";
SPrimeVec::usage="Return the spin-1 operators for a given directin.";
EigenState::usage="A certain eigenstates of spin-1 operators along the given direction.";

SingleNV::usage="Hamiltonian of a single NV";

EspinDipoleIntSecular::usage="Secular Hamiltonian of dipolar coupling of two electron spins";
NVPair::usage="Hamiltonian of a pair of NVs";

Observable::usage="mean value of an operator on a given state";


Begin["`Private`"] (* Begin Private Context *) 

transferMatrix[zdir_]:=
	Module[{res, xdir, ydir},
	xdir = Normalize[Cross[{1., 0., 0.}, zdir]];
	ydir = Cross[zdir, xdir];
	res = { 
		{xdir.$Xdir, xdir.$Ydir, xdir.$Zdir},
		{ydir.$Xdir, ydir.$Ydir, ydir.$Zdir},
		{zdir.$Xdir, zdir.$Ydir, zdir.$Zdir}
		};
	
	Return[res];
	]

spinMatrixConvert[zdir_,{xComponent_,yComponent_, zComponent_}]:= 
	Module[{res, mat},
	mat = transferMatrix[zdir];
	res = mat.{xComponent,yComponent, zComponent};
	Return[res];
	]

SPrimeVec[zdir_]:=
	Module[{res},
		res=spinMatrixConvert[zdir,{SPIN1x, SPIN1y, SPIN1z}];
		Return[res];
	]

PrimeBasis[zdir_]:=
	Module[{res, szprime, u, v},
		szprime=SPrimeVec[zdir][[3]];
		{u,v}=Eigensystem[szprime];

		res=v[[Reverse[Ordering[u]]]];
		Return[res];
	]

EigenState[nvdir_,idx_]:=PrimeBasis[nvdir][[idx]]	

Observable[operator_,state_] := 
	Module[{res}, 
		res = Conjugate[state].operator.state;
		Return[res];
	]
	
SingleNV[nvdir_, b_, theta_, phi_, nuc_]:= 
	Module[{res, eIdentity, nIdentity, eHami, nHami, hfHami},
	
	eIdentity = IdentityMatrix[3];
	nIdentity = If[nuc=="N14", IdentityMatrix[3], IdentityMatrix[2]];
	
	eHami = NVelectronSpin[nvdir, b, theta, phi];
	nHami = If[nuc=="N14", NVn14Spin[nvdir, b, theta, phi], NVn15Spin[nvdir, b, theta, phi]];
	hfHami = If[nuc=="N14", NV14hf[nvdir], NV15hf[]];
	
	res = KroneckerProduct[eHami, nIdentity] + KroneckerProduct[eIdentity,nHami] + hfHami;
	Return[res];
	]/;(nuc=="N14" || nuc=="N15")

NVPair[{nv1dir_, nuc1_}, {nv2dir_, nuc2_}, Jval_, b_, theta_, phi_]:= 
	Module[{res, dim1, nv1, dim2, nv2, dip},
	dim1 = 3 * If[nuc1=="N14",3, 2];
	dim2 = 3 * If[nuc2=="N14",3, 2];
	nv1 = SingleNV[nv1dir, b, theta, phi, nuc1];
	nv2 = SingleNV[nv2dir, b, theta, phi, nuc2];
	dip = EspinDipoleIntSecular[Jval, nuc1, nuc2];
	
	res = KroneckerProduct[nv1, IdentityMatrix[dim2]]
		 + KroneckerProduct[IdentityMatrix[dim1], nv2]
		 + dip;
	Return[res];
	]

NVelectronSpin[nvdir_,b_,theta_,phi_] := 
	Module[{res,bx, by, bz, sxprime, syprime, szprime, H0, Hz},
	{sxprime, syprime, szprime} = SPrimeVec[nvdir];
	
	H0=$DiamondZFS*(szprime.szprime);
	H0=H0  * 10^9 / 1000 * 2 * Pi (*convert GHz to k rad/s*);
	
	{bx, by, bz} = b * {Sin[theta]Cos[phi], Sin[theta]Cos[phi], Cos[theta]};(*b in Tesla*)
	Hz = -$GammaE*(bx*SPIN1x+by*SPIN1y+bz*SPIN1z)/1000.(*unit: k rad/s*);
		
	res = H0+Hz;	
	Return[res];		
	]
	
NVn14Spin[nvdir_,b_,theta_,phi_] := 
	Module[{res, bx, by, bz, sxprime, syprime, szprime, H0, Hz},
	{sxprime, syprime, szprime} = SPrimeVec[nvdir];

	H0=$NVn14ZFS*(szprime.szprime);
	H0=H0  * 10^6 / 1000 * 2 * Pi (*convert MHz to k rad/s*);
	
	{bx, by, bz} = b * {Sin[theta]Cos[phi], Sin[theta]Cos[phi], Cos[theta]};(*b in Tesla*)
	Hz = -$GammaN14*(bx*SPIN1x+by*SPIN1y+bz*SPIN1z)/1000.(*unit: k rad/s*);
		
	res = H0+Hz;	
	Return[res];		
	]
	
NVn15Spin[nvdir_,b_,theta_,phi_] := 
	Module[{res, bx, by, bz, Hz},

	{bx, by, bz} = b * {Sin[theta]Cos[phi], Sin[theta]Cos[phi], Cos[theta]};(*b in Tesla*)
	Hz = -$GammaN15*(bx*SPINhalfx+by*SPINhalfy+bz*SPINhalfz)/1000.(*unit: k rad/s*);
		
	res = Hz;	
	Return[res];		
	]

NV14hf[nvdir_] := 
	Module[{res, sxprime, syprime, szprime, ixprime, iyprime, izprime},
	{sxprime, syprime, szprime} = {ixprime, iyprime, izprime} = SPrimeVec[nvdir];
	
	res =  ($HFn14vert * KroneckerProduct[sxprime, ixprime] +
			$HFn14vert * KroneckerProduct[syprime, iyprime] +
			$HFn14zz * KroneckerProduct[szprime, izprime] ) * 1000. * 2 * Pi (*convert MHz to k rad/s*); 
	Return[res];
	]	 

NV15hf[] := 
	Module[{res},
	res = $HFn15 * (KroneckerProduct[SPIN1x, SPINhalfx] 
					+ KroneckerProduct[SPIN1y, SPINhalfy] 
					+ KroneckerProduct[SPIN1z, SPINhalfz]
					)* 1000. * 2 * Pi (*convert MHz to k rad/s*); 
	Return[res];
	]	 

EspinDipoleIntSecular[Jval_, nuc1_, nuc2_] :=
	Module[{res, id1, id2, s1x, s1y, s1z, s2x, s2y, s2z},
	id1 = If[nuc1=="N14",IdentityMatrix[3], IdentityMatrix[2]];
	id2 = If[nuc2=="N14",IdentityMatrix[3], IdentityMatrix[2]];
	
	{s1x, s1y, s1z} = (KroneckerProduct[#,id1]&) /@ {SPIN1x, SPIN1y, SPIN1z};
	{s2x, s2y, s2z} = (KroneckerProduct[#,id2]&) /@ {SPIN1x, SPIN1y, SPIN1z};
	
	res = Jval * (-KroneckerProduct[s1x, s2x] - KroneckerProduct[s1y, s2y] + 2.0 * KroneckerProduct[s1z, s2z]);
	Return[res];
	]/;( (nuc1=="N14" || nuc1=="N15") && (nuc2=="N14" || nuc2=="N15") )
	
End[] (* End Private Context *)

EndPackage[]